﻿#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys
import os
import time
import signal

from urlparse import urlparse
from optparse import OptionParser

from PyQt4.QtCore import *
from PyQt4.QtGui import *
from PyQt4.QtWebKit import QWebView, QWebSettings

class QwkRenderer(QObject):
	def __init__(self):
		
		#initialize parameters
		self.url		=	None
		self.width		=	0
		self.height		=	0
		self.timeout	=	0
		self.ratio		=	Qt.IgnoreAspectRatio
		self.swidth		=	0
		self.sheight	=	0
		self.output		=	None
		self.format		=	0
		self.quality	=	-1

		#create a Qt Webkit viewer
		self.browser = QWebView()
		
		#binding state event
		self.connect(self.browser, SIGNAL("loadFinished(bool)"), self.__onLoadFinished)
		self.connect(self.browser, SIGNAL("loadStarted()"), self.__onLoadStarted)

		#config default options of browser
		self.browser.settings().setAttribute(QWebSettings.JavascriptEnabled, True)
		self.browser.settings().setAttribute(QWebSettings.PluginsEnabled, True)
		self.browser.settings().setAttribute(QWebSettings.PrivateBrowsingEnabled, True)
		self.browser.settings().setAttribute(QWebSettings.JavascriptCanOpenWindows, False)

		#initialize state
		self.__loadState = False
		self.__loadResult = False

	def wkRender(self):

		if self.url == None:
			raise RuntimeError("Request URL not given!")

		stopTime = time.time() + self.timeout
		
		self.browser.load(QUrl(self.url))
		
		while self.__loadState:
			if self.timeout > 0 and time.time() >= stopTime:
				raise RuntimeError("Request timed out!")
			QCoreApplication.processEvents()

		if self.__loadResult == False:
			raise RuntimeError("Failed to load %s" % self.url)
		
		#get webpage from browser
		webpage = self.browser.page()

		#delete scrollbar
		webpage.mainFrame().setScrollBarPolicy(Qt.Horizontal, Qt.ScrollBarAlwaysOff)
		webpage.mainFrame().setScrollBarPolicy(Qt.Vertical, Qt.ScrollBarAlwaysOff)

		#reset webpage size
		size = webpage.mainFrame().contentsSize()
		if self.width > 0:
			size.setWidth(self.width)
		if self.height > 0:
			size.setHeight(self.height)
		webpage.setViewportSize(size)

		#generate an image
		image = QImage(webpage.viewportSize(), QImage.Format_ARGB32)
		painter = QPainter(image)
		webpage.mainFrame().render(painter)
		painter.end()
		self.__saveAsImage(image)

	def __onLoadStarted(self):
		self.__loadState = True

	def __onLoadFinished(self, res):
		self.__loadState = False
		self.__loadResult = res

	def __saveAsImage(self, image):
		if self.swidth > 0 and self.sheight > 0:
			image = image.scaled(self.swidth, self.sheight, self.ratio)
		
		name, suffix = os.path.splitext(self.output)
		
		if not suffix:
			if self.format:
				name += '.' + self.format
			else:
				name += '.png'
		else:
			if self.format and self.format not in suffix:
				name = self.output + '.' + self.format
			else:
				name = self.output
		
		if image.save(name, self.format, self.quality) == False:
			raise RuntimeError("Image save error!")
		
	
	#set parameter methods
	def setRequestURL(self, url):
		if url:
			if not urlparse(url).scheme:
				url = 'http://' + url
		self.url = url

	def setScreenSize(self, size):
		self.width, self.height = size
	
	def setTimeOut(self, time):
		self.timeout = time

	def setAspectRatio(self, ratio):
		if ratio == 'keep':
			self.ratio = Qt.KeepAspectRatio
		elif ratio == 'expand':
			self.ratio = Qt.KeepAspectRatioByExpanding
		else:
			self.ratio = Qt.IgnoreAspectRatio
	
	def setScaleSize(self, size):
		self.swidth, self.sheight = size

	def setOutputFile(self, output):
		self.output = output

	def setImageFormat(self, format):
		if format == 0:
			format = None
		self.format = format

	def setImageQuality(self, quality):
		self.quality = quality

if __name__ == '__main__':

	#Parse command line arguments and validate them
	description = '''Creates a screenshot of a website using QtWebKit.
					 This is free software, and you are welcome to redistribute.
					 it under the terms of the GNU GPL v2.'''

	parser = OptionParser(usage = "usage: %prog [options] <URL> <OUTFILE>",
						  version = "%prog 0.1, Copyright (c) 2012 Lianming Du",
						  description = description)
	
	parser.add_option("-s", "--screen", dest="screen", nargs=2, default = (0, 0), type = "int", help="Set screen height and width (default is the size of web page content) [default: %default].", metavar = "WIDTH HEIGHT")
	parser.add_option("-c", "--scale", dest="scale", nargs=2, default = (0, 0), type="int", help="Scale the image to this size", metavar="WIDTH HEIGHT")
	parser.add_option("-r", "--ratio", dest="ratio", type="choice", default="ignore", choices=["ignore", "keep", "expand"], help="Set the scale ratio, the value is one of 'ignore', 'keep' or 'expand' [default: %default]")
	parser.add_option("-t", "--timeout", dest="timeout", default=0, type="int", help="The request timeout [default: %default]", metavar="SECONDS")
	parser.add_option("-f", "--format", dest="format", default=0, help="Output image format [default: %default] 0 means guess the format by looking at output filename's suffix", metavar="FORMAT")
	parser.add_option("-q", "--quality", dest="quality", default=-1, type="int", help="Set the qaulity of output image (range 0 to 100 or -1) [default: %default]")


	options, args = parser.parse_args()

	if len(args) != 2:
		parser.error("incorrect number of arguments")
	
	options.url, options.output = args
	
	app = QApplication(sys.argv)
	signal.signal(signal.SIGINT, signal.SIG_DFL)

	browser = QwkRenderer()
	
	#setting parameters
	browser.setRequestURL(options.url)
	browser.setScreenSize(options.screen)
	browser.setScaleSize(options.scale)
	browser.setAspectRatio(options.ratio)
	browser.setImageFormat(options.format)
	browser.setImageQuality(options.quality)
	browser.setOutputFile(options.output)

	#start render page to image
	def __onProcess():
		try:
			browser.wkRender()
			sys.exit(0)
		except RuntimeError, e:
			print e.message
			sys.exit(1)

	QTimer().singleShot(0, __onProcess)
	sys.exit(app.exec_())